﻿#include <iostream>
#include <fstream>
#include <cstring>
#include <map>
#include <queue>
#include <winsock2.h>
#include <windows.h>
#include "ikcp.h"
#include <chrono>
#include "receiver.h"

#define KCP_MTU 1400
#define LOCAL_PORT 5001
#define REMOTE_IP "127.0.0.1"
#define REMOTE_PORT 5000
#define MAGIC_LEN 4
#define CONTROL_MAGIC "CTRL"
#define DATA_MAGIC "DATA"
#define ACK_MAGIC "ACK_"

static inline uint32_t ikcp_getms() {
    return GetTickCount();
}

// 声明函数
DWORD WINAPI processConv(LPVOID lpParam);
DWORD WINAPI kcp_update_thread(LPVOID lpParam);

struct ConvData {
    std::queue<std::pair<char*, int>> queue;
    CRITICAL_SECTION cs;
    HANDLE hEvent;
    std::chrono::steady_clock::time_point createTime; // 记录 KCP 实例创建时间
    bool fileTransferStarted; // 标记文件传输是否开始
};

std::map<uint32_t, ConvData> convQueue;
std::map<uint32_t, ikcpcb*> convKCP;
std::map<uint32_t, std::ofstream> convFiles;
CRITICAL_SECTION mapMutex;
HANDLE hReceiveThread;
HANDLE hCleanupThread;
bool running = true;

int udp_output(const char *buf, int len, ikcpcb *kcp, void *user) {
    int sockfd = *(int*)user;
    struct sockaddr_in dest_addr;
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_port = htons(REMOTE_PORT);
    dest_addr.sin_addr.s_addr = inet_addr(REMOTE_IP);
    int ret = sendto(sockfd, buf, len, 0, (struct sockaddr*)&dest_addr, sizeof(dest_addr));
    if (ret < 0) {
        printf("[UDP] 发送失败: %d\n", WSAGetLastError());
    } else {
        printf("[UDP] 发送 %d 字节\n", ret);
    }
    return ret;
}

void sendAck(ikcpcb* kcp, uint32_t conv) {
    char ack_buf[MAGIC_LEN + 4];
    memcpy(ack_buf, ACK_MAGIC, MAGIC_LEN);
    uint32_t conv_net = htonl(conv);
    memcpy(ack_buf + MAGIC_LEN, &conv_net, 4);
    ikcp_send(kcp, ack_buf, MAGIC_LEN + 4);
    ikcp_update(kcp, ikcp_getms());
    printf("[ACK] 发送 ACK 回执, conv=%u\n", conv);
}

DWORD WINAPI receiveUDP(LPVOID lpParam) {
    int sockfd = *(int*)lpParam;
    char buffer[KCP_MTU + 16];
    struct sockaddr_in src_addr;
    int addr_len = sizeof(src_addr);

    while (running) {
        int recv_len = recvfrom(sockfd, buffer, KCP_MTU + 16, 0, (struct sockaddr*)&src_addr, &addr_len);
        if (recv_len > 0) {
            // 解析 conv 值（前 4 个字节）
            uint32_t conv = *(uint32_t*)(buffer);

            EnterCriticalSection(&mapMutex);

            // 判断是否存在对应的 KCP 实例
            if (convKCP.find(conv) == convKCP.end()) {
                // 创建新的 KCP 实例
                ikcpcb* kcp = ikcp_create(conv, (void*)&sockfd);
                kcp->output = udp_output;
                ikcp_wndsize(kcp, 256, 256); // 增大窗口大小
                ikcp_nodelay(kcp, 1, 10, 2, 1); // 启用无延迟模式
                convKCP[conv] = kcp;

                // 创建 ConvData 结构
                ConvData data;
                InitializeCriticalSection(&data.cs);
                data.hEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
                data.createTime = std::chrono::steady_clock::now(); // 记录创建时间
                data.fileTransferStarted = false; // 文件传输未开始
                convQueue[conv] = data;

                // 创建处理线程
                CreateThread(NULL, 0, processConv, (LPVOID)conv, 0, NULL);

                // 创建 KCP 更新线程
                CreateThread(NULL, 0, kcp_update_thread, (LPVOID)kcp, 0, NULL);

                printf("[KCP] 创建 KCP 实例, conv=%u\n", conv);
            }

            // 发送允许传文件的回执信息
            sendAck(convKCP[conv], conv);

            // 将数据包输入 KCP
            ikcpcb* kcp = convKCP[conv];
            int ret = ikcp_input(kcp, buffer, recv_len);
            if (ret < 0) {
                printf("[KCP] 输入数据失败, conv=%u, ret=%d\n", conv, ret);
            } else {
                printf("[KCP] 成功输入数据包, conv=%u, 长度=%d\n", conv, recv_len);
            }

            LeaveCriticalSection(&mapMutex);
        } else {
            printf("[UDP] 接收失败: %d\n", WSAGetLastError());
        }
    }
    return 0;
}

DWORD WINAPI processConv(LPVOID lpParam) {
    uint32_t conv = (uint32_t)lpParam;
    ikcpcb* kcp = convKCP[conv];
    std::ofstream& file = convFiles[conv];

    while (running) {
        WaitForSingleObject(convQueue[conv].hEvent, INFINITE);
        EnterCriticalSection(&convQueue[conv].cs);
        if (!convQueue[conv].queue.empty()) {
            std::pair<char*, int> packet = convQueue[conv].queue.front();
            convQueue[conv].queue.pop();
            LeaveCriticalSection(&convQueue[conv].cs);

            // 从 KCP 读取应用数据
            char app_data[1024];
            int recv_len = ikcp_recv(kcp, app_data, sizeof(app_data));
            if (recv_len > 0) {
                // 解析 magic 值
                char magic[MAGIC_LEN + 1];
                memcpy(magic, app_data, MAGIC_LEN);
                magic[MAGIC_LEN] = '\0';

                if (strcmp(magic, CONTROL_MAGIC) == 0) {
                    // 处理控制包
                    uint16_t name_len = ntohs(*(uint16_t*)(app_data + MAGIC_LEN));
                    std::string filename(app_data + MAGIC_LEN + 2, name_len);
                    printf("[控制] 收到控制包, conv=%u, 文件名=%s\n", conv, filename.c_str());

                    // 创建文件
                    std::ofstream file("recv_conv_" + std::to_string(conv) + "_" + filename, std::ios::binary);
                    if (file.is_open()) {
                        convFiles[conv] = std::move(file);
                        printf("[文件] 创建文件, conv=%u, 文件名=%s\n", conv, filename.c_str());
                    } else {
                        printf("[文件] 创建文件失败, conv=%u, 文件名=%s\n", conv, filename.c_str());
                    }

                    // 标记文件传输开始
                    convQueue[conv].fileTransferStarted = true;
                } else if (strcmp(magic, DATA_MAGIC) == 0) {
                    // 处理数据包
                    int data_len = recv_len - MAGIC_LEN;
                    file.write(app_data + MAGIC_LEN, data_len);
                    file.flush(); // 刷新文件缓冲区
                    printf("[数据] 接收数据包, conv=%u, 长度=%d\n", conv, data_len);
                } else {
                    printf("[数据] 未知的 magic: %s\n", magic);
                }
            } else if (recv_len < 0) {
                printf("[KCP] 接收数据错误, conv=%u, ret=%d\n", conv, recv_len);
            }
        } else {
            LeaveCriticalSection(&convQueue[conv].cs);
        }
    }

    // 释放资源
    file.close();
    ikcp_release(kcp);
    DeleteCriticalSection(&convQueue[conv].cs);
    CloseHandle(convQueue[conv].hEvent);
    convQueue.erase(conv);
    convKCP.erase(conv);
    convFiles.erase(conv);
    printf("[KCP] 释放 KCP 实例, conv=%u\n", conv);

    return 0;
}

DWORD WINAPI kcp_update_thread(LPVOID lpParam) {
    ikcpcb* kcp = (ikcpcb*)lpParam;
    while (running) {
        ikcp_update(kcp, ikcp_getms());
        Sleep(10); // 每 10ms 更新一次
    }
    return 0;
}

DWORD WINAPI cleanupKCP(LPVOID lpParam) {
    while (running) {
        EnterCriticalSection(&mapMutex);

        auto now = std::chrono::steady_clock::now();
        for (auto it = convKCP.begin(); it != convKCP.end(); ) {
            uint32_t conv = it->first;
            ConvData& data = convQueue[conv];

            // 如果 KCP 实例创建后 10 分钟依然没有文件开始传输，则删除该实例
            if (!data.fileTransferStarted && 
                std::chrono::duration_cast<std::chrono::minutes>(now - data.createTime).count() >= 10) {
                printf("[清理] 删除未使用的 KCP 实例, conv=%u\n", conv);

                // 释放资源
                ikcp_release(it->second);
                DeleteCriticalSection(&data.cs);
                CloseHandle(data.hEvent);
                convQueue.erase(conv);
                convFiles.erase(conv);
                it = convKCP.erase(it);
            } else {
                ++it;
            }
        }

        LeaveCriticalSection(&mapMutex);

        // 每隔 1 分钟检查一次
        Sleep(60000);
    }
    return 0;
}

int main() {
    // 设置控制台编码为 UTF-8
    SetConsoleOutputCP(65001);

    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        printf("初始化 Winsock 失败\n");
        return 1;
    }

    int sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sockfd == INVALID_SOCKET) {
        printf("创建套接字失败!\n");
        WSACleanup();
        return 1;
    }

    struct sockaddr_in local_addr;
    local_addr.sin_family = AF_INET;
    local_addr.sin_port = htons(LOCAL_PORT);
    local_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    bind(sockfd, (struct sockaddr*)&local_addr, sizeof(local_addr));

    InitializeCriticalSection(&mapMutex);
    hReceiveThread = CreateThread(NULL, 0, receiveUDP, (LPVOID)&sockfd, 0, NULL);
    hCleanupThread = CreateThread(NULL, 0, cleanupKCP, NULL, 0, NULL);

    // 等待用户输入以停止程序
    printf("按 Enter 键停止...\n");
    std::cin.get();
    running = false;

    WaitForSingleObject(hReceiveThread, INFINITE);
    WaitForSingleObject(hCleanupThread, INFINITE);
    CloseHandle(hReceiveThread);
    CloseHandle(hCleanupThread);
    closesocket(sockfd);
    WSACleanup();
    return 0;
}